package com.bigData.main.controller;

import com.bigData.main.dto.OllamaRequest;
import com.bigData.main.dto.QuestionRequest;
import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.databind.MappingIterator;
import com.fasterxml.jackson.databind.ObjectMapper;
import net.minidev.json.JSONObject;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.converter.StringHttpMessageConverter;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.Map;

@RestController
public class AIController {

    private static final String OLLAMA_API_URL = "http://localhost:11434/api/generate";

    @PostMapping("/api/ask")
    public ResponseEntity<String> askQuestion(@RequestBody QuestionRequest request) {
        RestTemplate restTemplate = new RestTemplate();
        restTemplate.getMessageConverters().add(0, new StringHttpMessageConverter(StandardCharsets.UTF_8));

        // 构造Ollama API请求
        OllamaRequest ollamaRequest = new OllamaRequest();
        ollamaRequest.setModel("deepseek-r1:14b");
        ollamaRequest.setPrompt(request.getQuestion());

        // 调用Ollama API并接收原始响应
        String rawResponse = restTemplate.postForObject(OLLAMA_API_URL, ollamaRequest, String.class);

        // 使用Jackson解析NDJSON流式响应
        StringBuilder answerBuilder = new StringBuilder();
        JsonFactory factory = new JsonFactory();
        ObjectMapper objectMapper = new ObjectMapper(factory);

        try (InputStream inputStream = new ByteArrayInputStream(rawResponse.getBytes(StandardCharsets.UTF_8))) {
            MappingIterator<Map> iterator = objectMapper.readerFor(Map.class)
                    .readValues(inputStream);

            while (iterator.hasNext()) {
                Map<String, Object> jsonMap = iterator.next();
                if (jsonMap.containsKey("response")) {
                    answerBuilder.append(jsonMap.get("response"));
                }
            }
        } catch (IOException e) {
            // 处理IO异常
            e.printStackTrace();
        }

        JSONObject jsonResponse = new JSONObject();
        jsonResponse.put("answer", answerBuilder.toString());
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_JSON_UTF8)
                .body(jsonResponse.toString());
    }
}
